import tensorflow as tf
import keras as keras
import numpy as np
import cv2
import matplotlib.pyplot as plt

print(keras.__version__)

(x_train, y_train), (x_valid, y_valid) = keras.datasets.mnist.load_data()
assert x_train.shape == (60000, 28, 28)
assert x_valid.shape == (10000, 28, 28)
assert y_train.shape == (60000,)
assert y_valid.shape == (10000,)
print("y_valid type is %s" % (y_valid.shape))
# step1: use sequential
model = keras.models.Sequential()

# step2: add layer
model.add(keras.layers.Flatten(input_shape=(x_train.shape[1], x_train.shape[2])))
model.add(keras.layers.Dense(units=784, activation="relu", input_dim=784))
model.add(keras.layers.Dense(units=10, activation="softmax"))

# step3: compile model
model.compile(optimizer="Adam", loss='sparse_categorical_crossentropy', metrics=['accuracy'])

print("model:")
model.summary()

# step4: train
model.fit(x_train, y_train, batch_size=64, epochs=10)

# step5: evaluate model
model.evaluate(x_valid, y_valid)

# save model
# model.save('keras_mnist.h5')

# img = x_valid[0]
img = cv2.imread("44.png")  # 读取图片
img = cv2.resize(img, (28, 28))  # 重置图片大小
img = img[:, :, 0:1]
plt.imshow(img[:, :, ::-1])

img = np.reshape(img, (-1, 28, 28))
# shape==== (1, 28, 28)
print('shape====', img.shape)
output = model.predict(img)
print("output type is %s" % (type(output)))
print(output)
predict_num = np.argmax(output, axis=1)  # 需要使用np.argmax找到最大值
print("predict num is ", predict_num)
